#!/usr/bin/python3
#coding:utf-8
import serial
import time
import threading
import numpy as np
import struct
import os
import socket
import cv2
import numpy as np

color_dist = {'red': {'Lower': np.array([0, 60, 60]), 'Upper': np.array([6, 255, 255])},
              'green': {'Lower': np.array([35, 43, 35]), 'Upper': np.array([90, 255, 255])},
              }


class NANO:
    def __init__(self):
        self.serial_path = "/dev/ttyUSB0"
        self.serial_rate = 115200
        
        self.init_vision()
        self.vision_run = True
        self.vision_run = True
        self.vision_thread = threading.Thread(target = self.vision_detect, args=())

        self.init_serial()
        self.serial_run = True
        self.serial_thread = threading.Thread(target = self.receive_arduino_func, args=())

        self.init_udp()
        self.udp_run = True
        self.udp_status = 0
        self.udp_thread = threading.Thread(target = self.receive_udp_func, args=())
    
    def init_vision(self):
        self.video_index = 0
        self.cap = cv2.VideoCapture(0)
        self.square_thresh = 50
        self.V_thresh = 100
        self.jpegQuality = 20


    def vision_detect(self):
        while self.vision_run:
            # start_time = time.time()
            _,img = self.cap.read()   #x/2 == 320
            # img_src = img.copy()
            img = cv2.cvtColor(img,cv2.COLOR_BGR2HSV)
            img_hsv = img.copy()
            # img = cv2.erode(img, None, iterations=2)

            color_list = ['red','green']
            for each_color in color_list:
                inrange_img = cv2.inRange(img, color_dist[each_color]['Lower'], color_dist[each_color]['Upper'])
                contours = cv2.findContours(inrange_img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[-2]
                bounding_boxes = [cv2.boundingRect(cnt) for cnt in contours]
                center_x_list = []
                center_y_list = []
                HSV_list = []
                for bbox in bounding_boxes:
                    [x , y, w, h] = bbox
                    if w*h < self.square_thresh:
                        continue
                    center_x = int(x + w / 2)
                    center_y = int(y + h / 2)
                    center_x_list.append(center_x)
                    center_y_list.append(center_y)
                    HSV_list.append(img_hsv[center_y][center_x][2])
                    # img_src = cv2.rectangle(img_src, (x, y), (x + w, y + h), (255, 0, 0), 2)
                if len(HSV_list) <= 0:
                    self.send_led_info(1000,0 if each_color == 'red' else 1)
                    continue
                target = np.argsort(HSV_list)[-1]
                if HSV_list[target] < self.V_thresh:
                    self.send_led_info(1000,0 if each_color == 'red' else 1)
                    continue
                target_x = center_x_list[target]
                target_y = center_y_list[target]
                # img_src = cv2.circle(img_src,(target_x,target_y),radius=10,color = (0,0,255) if each_color == 'red' else (0,255,0),thickness=5)
                self.send_led_info(target_x,0 if each_color == 'red' else 1)
            
            # print("fps:",1/(time.time()-start_time))

            # img_src = cv2.resize(img_src, (640, 480))
            # data = cv2.imencode('.jpg', img_src, (cv2.IMWRITE_JPEG_QUALITY, self.jpegQuality))[1].tobytes()
            # self.send_udp(data)
            time.sleep(1/1000)

    def init_udp(self):
        s = socket.socket(socket.AF_INET,socket.SOCK_DGRAM)
        s.connect(('192.168.8.107',8888))
        self.ip = s.getsockname()[0]
        s.close()
        self.vision_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.vision_addr = (self.ip,8888) # 必须绑定自己电脑IP和port
        self.vision_socket.bind(self.vision_addr)
        self.targetDict = {}
        self.targetDict[(self.ip, 8888)] = 1

    def send_udp(self,data):
        for target in self.targetDict.keys():
            if target != self.vision_addr:
                self.vision_socket.sendto(data, target)

    def init_serial(self):
        times = 0
        while(1):
            try:
                if times >= 2:
                    print("failed to connect to arduino")
                    self.arduino_ser = None
                    break
                self.arduino_ser = serial.Serial(self.serial_path, self.serial_rate,timeout=1)
                print("success connect to arduino")
                break
            except:
                times += 1
                print("arduino serial time out")
            time.sleep(1)
    
    def receive_udp_func(self):
        while self.udp_run:
            res,addr = self.vision_socket.recvfrom(640*480)
            # res = res.decode("utf8")
            if not res:
                continue
            if self.targetDict.get(addr) == None:
                self.targetDict[addr] = 1
                print(addr,"connected")
            time.sleep(1/50)

    def send_arduino(self,string):
        try:
            self.arduino_ser.write(string)
        except:
            return
    
    def receive_arduino_func(self):
        if self.arduino_ser != None and self.arduino_ser.isOpen():
            try:
                res = self.arduino_ser.read(1)
                res = res.decode("utf8")
                if res != '':
                    print(res)
            except:
                print("failed to read plc")
                return
        time.sleep(1/50)

    def send_led_info(self,data_x,class_id):

        data_x = struct.pack('=I',data_x)
        class_id = struct.pack('=I',class_id)


        cmd_string = b''
        cmd_string += bytes([0xFF])
        for i in class_id:
            cmd_string += bytes([i])
        for i in data_x:
            cmd_string += bytes([i])
        self.send_arduino(cmd_string)


    def start(self):
        self.vision_thread.start()
        self.serial_thread.start()
        self.udp_thread.start()
            

if __name__ == '__main__':
    node1 = NANO()
    node1.start()
